#include <string.h>
#include <stdlib.h>

#include "queue.h"

#include "psd_tpool.h"
#include "psd_backend_slabs.h"
#include "psd_listener.h"
#include "psd_settings.h"
#include "psd_conn.h"
#include "psd_logger.h"
#include "psd_protocols.h"
#include "psd_config.h"
#include "psd_auth.h"
#include "psd_event.h"
#include "psd_xfer.h"
#include "ring_buffer.h"

#include "compat.h"

#define UDP_ENABLE 1
#define UDP_DGRAM (1024*64-40)

pthread_mutex_t xfer_total_lock;
pthread_cond_t xfer_total_cond;
uint64_t xfer_total_bytes;

int udp_sd;
struct sockaddr_in si_udp;
static psdSess *slab_sess;
static int psd_slabs_data_conn_handler(psdListener *listener, psdConn *conn, void *arg);
void *psd_slabs_handle_data_conn(void *arg);
void *psd_slabs_data_xfer_thread(void *arg);

unsigned long psd_slabs_send_xfer_conn(psdConn *conn, psdXferArgs *arg, xspSlabRec *rec);
unsigned long psd_slabs_recv_xfer_conn(psdConn *conn, psdXferArgs *arg, xspSlabRec *rec);

int psd_backend_slabs_start() {
	char *str_val;
	char *transfer_backend = NULL;

	if (psd_depot_settings_get("generic", "transfer_backend", &str_val) == 0) {
                transfer_backend = str_val;
        }
	else {
		psd_info(0, "skipping SLaBS backend startup");
		return 0;
	}
	
	if (strcmp("slabs-ingress", transfer_backend) == 0) {
		pthread_t slab_xfer_thread;
		
		slab_sess = psd_alloc_sess();
		if (!slab_sess) {
			psd_err(0, "could not allocate new slab session");
			goto error_exit;
		}
		
		gen_rand_hex(slab_sess->id, 2*XSP_SESSIONID_LEN+1);
		psd_info(5, "new SLaBS session with id %s", slab_sess->id);
		
		LIST_INIT(&slab_sess->parent_conns);
		LIST_INIT(&slab_sess->child_conns);
		LIST_INIT(&slab_sess->parent_data_conns);
		LIST_INIT(&slab_sess->child_data_conns);
		
		// open control and data connections to the next slabs GW
		// we make this GW explicit in the config for now
		// eventually we want to aggregate connections and find next slabs GW automatically
		
		// just let each connection setup its own control channel (usual session setup) for now
		//if (psd_slabs_connect_control(slab_sess) != 0) {
		//	psd_err(0, "couldn't start control connection to next SLaBS GW");
		//}
		
		// but we create a dedicated data channel and have threads on each side process the buffers
		if (psd_slabs_connect_data(slab_sess) != 0) {
			psd_err(0, "couldn't start data connection to next SLaBS GW");
			goto error_exit;
		}

		// start a thread that creates and sends slabs over this GW-GW session
		if (pthread_create(&slab_xfer_thread, NULL, psd_slabs_data_xfer_thread, (void *) slab_sess) < 0) {
			psd_err(0, "couldn't start slabs_data_xfer_thread");
                        goto error_exit;
                }
	}
	else if (strcmp("slabs-egress", transfer_backend) == 0) {

		// create a data listener
		// the handler starts a thread that receives slab buffer and notifies egress threads
		if (psd_slabs_start_data_listener(slab_sess) != 0) {
			psd_err(0, "couldn't start data listener");
			goto error_exit;
		}		 		
	}
	
	return 0;

 error_exit:
        return -1;
}

// need to refactor this thing
void *psd_slabs_data_xfer_thread(void *args) {
	psdXferArgs **xfer_args;
	psdXferArgs *curr_arg;
	psdConn *data_conn;
	psdSess *sess;
	xspMsg *msg;
	xspSlabRec *curr_rec;

	int xfer_count;
	int i;

	unsigned long buf_bytes;
	unsigned long bytes_sent;

	sess = (psdSess*) args;
	data_conn = LIST_FIRST(&sess->child_data_conns);
	if (!data_conn) {
		psd_err(0, "couldn't get slabs data conn!");
		return;
	}

	bytes_sent = 0;

	while (1) {
		pthread_mutex_lock(&xfer_total_lock);
		{
			if (xfer_total_bytes < PSD_SLAB_SIZE) {
				psd_info(5, "waiting to form a full slab");
				pthread_cond_wait(&xfer_total_cond, &xfer_total_lock);
			}
		}
		pthread_mutex_unlock(&xfer_total_lock);
		
		
		xfer_args = psd_get_xfer_args(&xfer_count);

		psd_info(10, "there are %d active xfers", xfer_count);
		psd_info(10, "total xfer buf size: %llu", xfer_total_bytes);
		
		// create and send the slab records
		unsigned int t_length = 0;
		
		xspSlabInfo *new_info = malloc(sizeof(xspSlabInfo));
		new_info->entries = malloc(xfer_count * sizeof(xspSlabRec*));
		new_info->rec_count = xfer_count;
		
		for (i=0; i<xfer_count; i++) {
			xspSlabRec *new_rec = malloc(sizeof(xspSlabRec));

			curr_arg = xfer_args[i];
			
			strcpy(new_rec->sess_id, curr_arg->sess->id);
			
			buf_bytes = ring_buffer_count_bytes(curr_arg->buf);
			
			new_rec->offset = t_length;
			new_rec->length = (uint32_t)((double)buf_bytes/(double)xfer_total_bytes*(double)PSD_SLAB_SIZE);
			// ignore checksum for now
			new_rec->crc = 0;
			
			t_length += new_rec->length;
			new_info->entries[i] = new_rec;
		}
		new_info->length = t_length;
		
		psd_info(10, "sending slab of length %lu for %d sessions", new_info->length, new_info->rec_count);

		if (psd_conn_send_msg(data_conn, XSP_MSG_SLAB_INFO, new_info) <= 0) {
			psd_err(0, "could not send slabs info, stopping!");
			return;
		}
		
		// now send the corresponding buffer data
		bytes_sent = 0;
		
		for (i=0; i<xfer_count; i++) {
			curr_arg = xfer_args[i];
			curr_rec = new_info->entries[i];

			pthread_mutex_lock(&(curr_arg->buf_lock));
			{
				if (UDP_ENABLE) {
					unsigned long xfer_amt = curr_rec->length;
					unsigned long sent = 0;
					int n;
					char *buf_ptr;
					
					unsigned long send_amt;
					
					psd_info(10, "id: %s, xfer_amt: %lu, slab size: %lu",
						 curr_arg->sess->id, xfer_amt, PSD_SLAB_SIZE);

					while (sent < xfer_amt) {

						if (UDP_DGRAM > (xfer_amt - sent))
							send_amt = (xfer_amt - sent);
						else
							send_amt = UDP_DGRAM;

						buf_ptr = ring_buffer_read_address(curr_arg->buf);
						n = sendto(udp_sd, buf_ptr, send_amt, 0, &(si_udp), sizeof(si_udp));
						
						psd_info(11, "sendto(%d) returned %d", udp_sd, n);
						if (n < 0) {
							perror("sendto: ");
							break;
						}

						ring_buffer_read_advance(curr_arg->buf, n);
						sent += n;
						//usleep(10);
					}
					bytes_sent += sent;
				}
				else {
					bytes_sent += psd_slabs_send_xfer_conn(data_conn, curr_arg, curr_rec);
				}
			}
			pthread_mutex_unlock(&(curr_arg->buf_lock));
		}

		/*
		psd_info(11, "waiting for SLAB ACK");
		msg = psd_conn_get_msg(data_conn, 0);
		if (!msg) {
			psd_err(0, "did not receive properly formed message, stopping!");
			return;
		}
		
		if (msg->type == XSP_MSG_SESS_ACK) {
			psd_info(10, "got ACK for last slab transfer");
		}
		*/
		
		pthread_mutex_lock(&xfer_total_lock);
                {
			psd_info(10, "removing %lu from xfer_total", bytes_sent);
                        xfer_total_bytes -= bytes_sent;
                }
                pthread_mutex_unlock(&xfer_total_lock);
		
		for (i=0; i<xfer_count; i++) {
                        curr_arg = xfer_args[i];
			if (ring_buffer_count_free_bytes(curr_arg->buf) > 0)
				pthread_cond_signal(&(curr_arg->buf_cond));
		}
	}
}

unsigned long psd_slabs_send_xfer_conn(psdConn *data_conn, psdXferArgs *curr_arg, xspSlabRec *curr_rec) {
	unsigned long conn_sent;
	unsigned long xfer_amt;
	unsigned long send_amt;
	char *buf_ptr;
	int n;

	xfer_amt = curr_rec->length;
	
	psd_info(10, "id: %s, xfer_amt: %lu, slab size: %lu",
		 curr_arg->sess->id, xfer_amt, PSD_SLAB_SIZE);
	
	// send to egress GW
	conn_sent = 0;
	while (conn_sent < xfer_amt) {
		
		if (PSD_BUF_SIZE < (xfer_amt - conn_sent)) {
			send_amt = PSD_BUF_SIZE;
		} else {
			send_amt = xfer_amt - conn_sent;
		}
		
		buf_ptr = ring_buffer_read_address(curr_arg->buf);
		n = psd_conn_write(data_conn, buf_ptr, send_amt, 0);
		ring_buffer_read_advance(curr_arg->buf, n);
		if (n < 0) {
			psd_err(0, "buffer send failed");
			break;
		}
		psd_info(11, "data conn send returned %d", n);
		conn_sent += n;
	}
	return conn_sent;
}

int psd_slabs_start_data_listener(psdSess *sess) {
	psdSettings *settings = NULL;
	psdListener *listener;
	char *protocol;
	char *default_protocol = "tcp";
	int port = 5020;

	if (psd_depot_settings_get_section("slabs", &settings) != 0) {
                psd_err(0, "No slabs config section found!");
                goto error_exit;
        }
	
	if (psd_settings_get_2(settings, "listeners", "protocol", &protocol) != 0) {
                psd_err(0, "No slabs data protocol found, using default %s", default_protocol);
		protocol = default_protocol;
        }

        if (psd_settings_get_int_2(settings, "listeners", "port", &port) != 0) {
                psd_err(0, "No slabs data port found, using default %d", port);
		port = 5020;
	}
	
	if (psd_settings_set_int_2(settings, protocol, "port", port) != 0) {
                psd_err(0, "couldn't set protocol port");
                goto error_exit_parsed;
        }

	psd_info(0, "Setting up listener for %s", protocol);

	if ((listener = psd_protocol_setup_listener(protocol, protocol, settings, 0, psd_slabs_data_conn_handler, NULL)) == NULL) {
		psd_err(0, "Couldn't setup listener for protocol %s", protocol);
		goto error_exit;
	}

	if (psd_listener_start(listener) != 0) {
		psd_err(0, "Couldn't start listener for protocol %s", protocol);
		goto error_exit;
	}

	return 0;

 error_exit_parsed:
	free(protocol);
 error_exit:
	psd_settings_free(settings);
	return -1;
}

static int psd_slabs_data_conn_handler(psdListener *listener, psdConn *conn, void *arg) {
        int retval;

        psd_info(0, "spawning slabs_handle_conn for %s", conn->description);

        retval = psd_tpool_exec(psd_slabs_handle_data_conn, conn);

        return retval;
}

void *psd_slabs_handle_data_conn(void *arg) {
	xspMsg *msg;
	psdConn *data_conn = (psdConn*) arg;
	int n;
	int i, j;
	int found = 0;

	psdXferArgs **xfer_args;
	psdXferArgs *curr_arg;
	int xfer_count;

	int recv_sd;

	psd_info(0, "handling new slabs data conn, waiting for slab info");
	
	if (UDP_ENABLE) {
		int buf_size = 4194304;
		
		recv_sd = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
		if (recv_sd < 0) {
			psd_err(0, "socket() failed");
		}
		
		if (setsockopt(recv_sd, SOL_SOCKET, SO_RCVBUF, &buf_size, sizeof(buf_size) != 0)) {
			psd_err(0, "couldn't set UDP recv buff");
			perror("setsockopt: ");
		}

		struct sockaddr_in si_me;
		memset((char *) &si_me, 0, sizeof(si_me));
		si_me.sin_family = AF_INET;
		si_me.sin_port = htons(5020);
		si_me.sin_addr.s_addr = htonl(INADDR_ANY);

		if (bind(recv_sd, &si_me, sizeof(si_me)) < 0){
			psd_err(0, "bind() failed");
		}
	}

	struct timeval start_time, end_time, prev_time;

	while (1) {
		
		msg = psd_conn_get_msg(data_conn, 0);
		if (!msg) {
			psd_err(5, "Did not receive properly formed message, stopping!");
                        return;
                }
		
		if (msg->type == XSP_MSG_SLAB_INFO) {
		        gettimeofday(&start_time, NULL);
			psd_info(10, "got SLAB_INFO message, processing...");
			
			xspSlabRec *curr_rec;
			xspSlabInfo *slab_info = (xspSlabInfo*) msg->msg_body;
			if (!slab_info) {
				psd_err(0, "couldn't get slab_info from message, stopping!");
				return;
			}
			
			psd_info(10, "receiving slab of size %lu for %d sessions", slab_info->length, slab_info->rec_count);

			xfer_args = psd_get_xfer_args(&xfer_count);
			psd_info(10, "there are %d active xfers", xfer_count);
			
			// test loop that just reads the slab and throws it away
			/*
			conn_recv = 0;
			char buf[PSD_BUF_SIZE];
			while (conn_recv < slab_info->length) {
				n = psd_conn_read(data_conn, buf, PSD_BUF_SIZE, 0);
				
				if (n <= 0)
					return;
				psd_info(11, "data conn recv returned %d", n);
				conn_recv += n;
			}
			*/
			
			for (i=0; i<slab_info->rec_count; i++) {
				curr_rec = slab_info->entries[i];
				
				psd_info(10, "id: %s, offset: %lu, length: %lu, crc: %lu", curr_rec->sess_id,
					 curr_rec->offset, curr_rec->length, curr_rec->crc);
				
				// find the right session connection to send the chunk to
				found = 0;
				for (j=0; j<xfer_count; j++) {
					curr_arg = xfer_args[j];
					if (!memcmp(curr_rec->sess_id, curr_arg->sess->id, 2*XSP_SESSIONID_LEN + 1)) {
						psd_info(11, "found xfer_arg");
						found = 1;
						break;
					}
				}
				if (!found) {
					psd_err(0, "couldn't find active session to send to, stopping!");
					return;
				}
				
				pthread_mutex_lock(&(curr_arg->buf_lock));
				{
					if (UDP_ENABLE) {
						unsigned long recvd;
						unsigned long buf_bytes;
						unsigned long recv_amt;
						char *buf_ptr;
						int n;
						
						struct sockaddr si_other;
						int slen;

						recvd = 0;
						while (recvd < curr_rec->length) {
							if (UDP_DGRAM > (curr_rec->length - recvd))
								recv_amt = (curr_rec->length - recvd);
							else
								recv_amt = UDP_DGRAM;
							
							// we are receiving faster than the egress connection can write
							// signal and yield to the writer thread
							while ((buf_bytes = ring_buffer_count_free_bytes(curr_arg->buf)) < recv_amt) {
								pthread_mutex_unlock(&(curr_arg->buf_lock));
								pthread_cond_signal(&(curr_arg->buf_cond));
								sched_yield();
								pthread_mutex_lock(&(curr_arg->buf_lock));
							}

							buf_ptr = ring_buffer_write_address(curr_arg->buf);
							n = recvfrom(recv_sd, buf_ptr, recv_amt, 0, &si_other, &slen);
							psd_info(11, "recvfrom(%d) returned %d", recv_sd, n);
							if (n < 0) {
								perror("recvfrom: ");
								break;
							}

							ring_buffer_write_advance(curr_arg->buf, n);
							recvd += n;
						}
					}
					else {
						psd_slabs_recv_xfer_conn(data_conn, curr_arg, curr_rec);
					}
				}
				pthread_mutex_unlock(&(curr_arg->buf_lock));
			}

			// don't sent ack, let's just print how long it took to get the slab data
			gettimeofday(&end_time, NULL);
			
			printf("%.2f, %.2f, %.2f, %.2f\n",
			       (double)slab_info->length/1e6,
			       (((double) slab_info->length*8)/difftv(&start_time, &end_time))/1e6,
			       difftv(&start_time, &end_time),
			       difftv(&prev_time, &end_time));
			
			gettimeofday(&prev_time, NULL);
			
			// tell all the waiting threads that they have new data to write
			for (i=0; i<xfer_count; i++) {
				if (ring_buffer_count_bytes(xfer_args[i]->buf) > 0)
					pthread_cond_signal(&(xfer_args[i]->buf_cond));
			}
			
			//psd_info(10, "sending SLAB ACK");
			// send back an ack for the slab we just got
			//if (psd_conn_send_msg(data_conn, XSP_MSG_SESS_ACK, NULL) <= 0) {
			//	psd_err(0, "could not send ACK, stopping!");
			//	return;
			//}
		}
		else {
			psd_err(0, "did not get a SLAB_INFO message, stopping!");
			return;
		}
	}
}

unsigned long psd_slabs_recv_xfer_conn(psdConn *data_conn, psdXferArgs *curr_arg, xspSlabRec *curr_rec) {
	unsigned long conn_recv;
	unsigned long buf_bytes;
        unsigned long recv_amt;
        unsigned long xfer_amt;
	char *buf_ptr;
	int n;

	conn_recv = 0;
	while (conn_recv < curr_rec->length) { 		

		if (PSD_BUF_SIZE < (curr_rec->length - conn_recv)) {
			recv_amt = PSD_BUF_SIZE;
		} else {
			recv_amt = curr_rec->length - conn_recv;
		}

		// we are receiving faster than the egress connection can write
		// signal and yield to the writer thread
		while ((buf_bytes = ring_buffer_count_free_bytes(curr_arg->buf)) < recv_amt) {
			pthread_mutex_unlock(&(curr_arg->buf_lock));
			pthread_cond_signal(&(curr_arg->buf_cond));
			sched_yield();
			pthread_mutex_lock(&(curr_arg->buf_lock));
		}

		buf_ptr = ring_buffer_write_address(curr_arg->buf);
		
		psd_info(11, "slab record %s going into recv (%lu) with %lu free slab bytes",
			 curr_rec->sess_id, recv_amt, buf_bytes);
		
		n = psd_conn_read(data_conn, buf_ptr, recv_amt, 0);
		
		if (n <= 0)
			return;

		psd_info(11, "data conn recv returned %d", n);
		
		ring_buffer_write_advance(curr_arg->buf, n);
		conn_recv += n;
	}
	return conn_recv;
}

int psd_slabs_connect_control(psdSess *sess) {
	psdSettings *settings = NULL;
	psdConn *new_conn;
	xspMsg *msg;
	char *hostname;
	int port = 5006;
	char *protocol = "tcp";

	if (psd_depot_settings_get_section("slabs", &settings) != 0) {
                psd_err(0, "No slabs config section found!");
                goto error_exit;
        }
	
	if (psd_settings_get(settings, "host", &hostname) != 0) {
		psd_err(0, "No slabs control host found");
		goto error_exit;
	}

	if (psd_settings_get_int_2(settings, "control", "port", &port) != 0) {
                psd_err(0, "No slabs control port found, using default %d", port);
                port = 5006;
        }

	if (psd_settings_set_int_2(settings, protocol, "port", port) != 0) {
                psd_err(0, "couldn't set protocol port");
                goto error_exit;
        }

	new_conn = psd_protocol_connect_host(hostname, protocol, settings);
	if (!new_conn) {
                psd_err(5, "connect() to \"%s\" using \"%s\" protocol failed", hostname, protocol);
                goto error_exit_parsed;
        }
	
	return 0;
		
 error_exit_msg:
        xsp_free_msg(msg);
 error_exit_conn:
        psd_conn_shutdown(new_conn, (PSD_RECV_SIDE | PSD_SEND_SIDE));
 error_exit_parsed:
        free(hostname);
 error_exit:
	psd_settings_free(settings);
        return -1;
}

int psd_slabs_connect_data(psdSess *sess) {
	psdSettings *settings = NULL;
        psdConn *new_conn;
        xspMsg *msg;
        char *hostname;
        int port = 5020;
        char *protocol;
	char *use_udp;
	char *default_protocol = "tcp";

        if (psd_depot_settings_get_section("slabs", &settings) != 0) {
                psd_err(0, "No slabs config section found!");
                goto error_exit;
        }

        if (psd_settings_get(settings, "host", &hostname) != 0) {
	        psd_err(0, "No slabs data host found");
	        goto error_exit;
        }
	
	if (psd_settings_get_2(settings, "data", "protocol", &protocol) != 0) {
                psd_err(0, "No slabs data protocol found, using default %d", default_protocol);
		protocol = default_protocol;
        }

        if (psd_settings_get_int_2(settings, "data", "port", &port) != 0) {
                psd_err(0, "No slabs data port found, using default %d", port);
		port = 5020;
        }

        if (psd_settings_set_int_2(settings, protocol, "port", port) != 0) {
                psd_err(0, "couldn't set protocol port");
                goto error_exit_parsed1;
        }

        new_conn = psd_protocol_connect_host(hostname, protocol, settings);
        if (!new_conn) {
                psd_err(5, "connect() to \"%s\" using \"%s\" protocol failed", hostname, protocol);
                goto error_exit_parsed1;
        }

	LIST_INSERT_HEAD(&(sess->child_data_conns), new_conn, sess_entries);
	
	if (UDP_ENABLE) {
		udp_sd = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
		if (udp_sd < 0) {
			psd_err(0, "socket() failed");
			goto error_exit;
		}
		
		memset((char *) &si_udp, 0, sizeof(si_udp));
		si_udp.sin_family = AF_INET;
		si_udp.sin_port = htons(port);
		
		if (inet_aton(hostname, &si_udp.sin_addr) == 0) {
			psd_err(0, "inet_aton() failed\n");
		}
	}
	
        return 0;

 error_exit_conn:
        psd_conn_shutdown(new_conn, (PSD_RECV_SIDE | PSD_SEND_SIDE));
 error_exit_parsed1:
        free(protocol);
 error_exit_parsed2:
	free(hostname);
 error_exit:
	psd_settings_free(settings);
        return -1;
}

psdSess *psd_slabs_get_sess() {
	return slab_sess;
}
